import numpy as np
from util import createXY  # 用于创建数据集的特征和标签
from sklearn.model_selection import train_test_split
from lazypredict.Supervised import LazyClassifier  # 引入LazyPredict
import logging
import joblib  # 用于保存模型

# 配置logging, 确保能够打印正在运行的函数名
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 主函数，运行训练过程
def main():
    logging.info("开始训练流程...")
    
<<<<<<< HEAD
    # 载入和预处理数据，只使用flat特征提取方法
    X, y = createXY(train_folder="D:/tuxiang/data/train", dest_folder=".", method='flat')
=======
    # 根据mode初始化FAISS所需的资源
    res = faiss.StandardGpuResources() if args.mode == 'gpu' else None
    logging.info(f"选择模式是 {args.mode.upper()}")
    logging.info(f"选择特征提取方法是 {args.feature.upper()}")
    logging.info(f"选择使用的库是 {args.library.upper()}")

    # 载入和预处理数据
    X, y = createXY(train_folder="D:/tuxiang/data/train", dest_folder=".",method=args.feature)
>>>>>>> 96218815913d4a5f3b67b9d7d4ff4ffd6e900666
    X = np.array(X)
    y = np.array(y)
    logging.info("数据加载和预处理完成。")

    # 数据集分割为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")

    # 使用LazyPredict进行多模型训练
    logging.info("使用LazyPredict进行模型训练...")
    clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
    models, predictions = clf.fit(X_train, X_test, y_train, y_test)

    # 输出模型结果
    logging.info("模型训练完成，结果如下：")
    print(models)

    # 找到准确率最高的模型
    best_model_name = models.index[0]
    best_accuracy = models.iloc[0]["Accuracy"]
    logging.info(f'准确率最高的模型是: {best_model_name}, 准确率: {best_accuracy}')

    # 保存最佳模型
    best_model = clf.models[best_model_name]
    joblib.dump(best_model, 'best_model.pkl')
    logging.info(f"最佳模型已保存到 'best_model.pkl'.")

# 如果是主脚本，则执行main函数
if __name__ == '__main__':
    main()
